# coding=utf-8
'''
Created on 2020-4-14
Rewrite on 2021-5-20
@author: Yoga, Kuan

分割训练样本
'''

import re

import numpy as np
from sklearn.model_selection import train_test_split

def clear_text(text):
    p = re.compile(r"[^\u4e00-\u9fa5^0-9^a-z^A-Z\-、，。！？：；（）《》【】,!\?:;[\]()]")  # 匹配不是中文、数字、字母、短横线的部分字符
    return p.sub('', text)  # 将text中匹配到的字符删除

def clear_texts(df_content):
    clear_content_list = []
    for i, content in enumerate(df_content):
        clear_content_list.append(clear_text(content))
    return clear_content_list

def write2file(df_data, path, label_list):
    df_data = df_data[df_data.label.isin(label_list)]
    print(df_data.shape)
    print(df_data['label'].value_counts())
    X_train_df, X_test_df, y_train_df, y_test_df = train_test_split(
        df_data, df_data['label'], test_size=0.1,
        random_state=1, stratify=df_data['label']
    )
    X_train_df, X_val_df, y_train_df, y_val_df = train_test_split(
        X_train_df, X_train_df['label'], test_size=0.2,
        random_state=1, stratify=X_train_df['label']
    )
    print('label_list:', label_list)
    print('训练集size:', X_train_df.shape, y_train_df.shape)
    print('验证集size:', X_val_df.shape, y_val_df.shape)
    print('测试集size:', X_test_df.shape, y_test_df.shape)
    X_train_df.to_csv(
        path + 'train.txt', 
        sep='\t', index=False, encoding='UTF-8'
    )
    X_val_df.to_csv(
        path + 'val.txt', 
        sep='\t', index=False, encoding='UTF-8'
    )
    X_test_df.to_csv(
        path + 'test.txt', 
        sep='\t', index=False, encoding='UTF-8'
    )

def main(sep_type, label_list):
    '''
    分割数据集, 0和7两个背景类默认保留
    有两种模式: type='bi', type='multi'
    sep_type='bi': 
        按照label_list中指定的label分割二分类数据集
        Usage:
            main(type='bi', label_list=[1,4,5,6])
    type='multi': 
        按照label_list中指定的label分割多分类数据集
        Usage:
            main(type='multi', label_list=[1,4,5,6])
    '''
    # 读取数据集, 打印前十样本为例, 以及label分布
    df_data = pd.read_csv(
        './source/data/js_pd_seg_merge.txt',
        encoding='UTF-8', 
        sep='\t', 
        header=0 # 首行表头
    )
    print(df_data[:10])
    print(df_data['label'].value_counts())
    if sep_type == 'bi':
        for label in label_list:
            tmp = [label]
            tmp.append(0)
            tmp.append(7)
            file_name = 'label' + str(label)
            write2file(df_data, './source/data/bi_classification/' + file_name + '/', tmp)
    elif sep_type == 'multi':
        label_list.append(0)
        label_list.append(7)
        write2file(df_data, './source/data/multi_classification/', label_list)
    else:
        print('注意sep_type参数!')
        input()
